"""
python fujian1_gridsearchAIC_pmdarima_single_category.py
"""
import os
import json
import pandas as pd
from pmdarima import ARIMA
import logging
import numpy as np
from pmdarima import auto_arima

def evaluate_arima_model(X, p, d, q):
    """
    使用 ARIMA 模型评估给定的参数组合并返回 AIC。
    
    参数:
    - X: array-like, 输入时间序列数据
    - p: int, 自回归项数
    - d: int, 差分次数
    - q: int, 移动平均项数
    
    返回:
    - float: 模型的 AIC 值
    """
    try:
        model = ARIMA(X, order=(p, d, q))
        model.fit()
        return model.aic()
    except:
        return np.inf  # 返回无穷大表示模型无法拟合

def grid_search_arima(X, p_values, d_values, q_values):
    """
    网格搜索 ARIMA 参数，返回最佳参数和对应的 AIC。
    
    参数:
    - X: array-like, 输入时间序列数据
    - p_values: list, 自回归项数范围
    - d_values: list, 差分次数范围
    - q_values: list, 移动平均项数范围
    
    返回:
    - tuple: (最佳 p, d, q, 最佳 AIC)
    """
    best_aic = np.inf
    best_pdq = (0, 0, 0)

    for p in p_values:
        for d in d_values:
            for q in q_values:
                aic = evaluate_arima_model(X, p, d, q)
                if aic < best_aic:
                    best_aic = aic
                    best_pdq = (p, d, q)
                print(f'Tried ARIMA({p}, {d}, {q}) - AIC: {aic}')

    return best_pdq + (best_aic,)

def predict_and_save_inventory(input_file_path, output_file_path, log_path=None):
    # 设置日志路径
    if log_path is None:
        log_path = os.getcwd()

    log_file_path = os.path.join(log_path, 'prediction.log')
    logging.basicConfig(filename=log_file_path, level=logging.WARNING, format='%(asctime)s - %(levelname)s - %(message)s')

    try:
        # 读取输入 JSON 文件
        with open(input_file_path, 'r') as f:
            data = json.load(f)

        # 将数据转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)

        # 确保有 12 个数据点进行训练
        if len(df) != 12:
            logging.warning("输入数据点不足 12 个，无法进行 ARIMA 预测。")
            return

        # 填充缺失值（如果有）
        df['inventory'] = df['inventory'].ffill()

        # 检查是否存在 NaN 值
        if df['inventory'].isnull().any():
            logging.error("数据包含 NaN 值，无法进行预测。")
            return
        
        print("输入给arima的数据：", df['inventory'])

        # 网格搜索参数范围
        p_values = range(1, 4)  # p值范围
        d_values = range(0, 3)  # d值范围
        q_values = range(1, 4)  # q值范围

        # 进行网格搜索
        best_p, best_d, best_q, best_aic = grid_search_arima(df['inventory'], p_values, d_values, q_values)

        print("最佳 ARIMA 模型参数 (p, d, q):", (best_p, best_d, best_q))
        print("最佳 AIC:", best_aic)

        # 使用最佳参数建立最终模型
        model = ARIMA(df['inventory'], order=(best_p, best_d, best_q))
        model.fit()

        # 预测未来 3 个数据点
        forecast, conf_int = model.predict(n_periods=3, return_conf_int=True)
        
        print("预测出来的三个月的结果：", forecast)

        # 生成未来日期
        forecast_dates = pd.date_range(start=df.index[-1] + pd.DateOffset(months=1), periods=3, freq='MS')

        # 创建预测结果 DataFrame
        forecast_df = pd.DataFrame({'date': forecast_dates, 'inventory': forecast})

        # 合并原始数据和预测数据
        result_df = pd.concat([df.reset_index(), forecast_df], ignore_index=True)

        # 将日期转换为字符串格式，以便 JSON 序列化
        result_df['date'] = result_df['date'].dt.strftime('%Y-%m-%dT00:00:00.000')

        # 转换为 JSON 格式并保存到输出文件
        result_json = result_df.to_dict(orient='records')
        with open(output_file_path, 'w') as f:
            json.dump(result_json, f, indent=4)

        print(f"预测结果已保存至 {output_file_path}")

    except Exception as e:
        logging.error(f"处理文件时发生错误：{e}")

# 示例用法
if __name__ == "__main__":
    input_file = r"..\fujian\fujian1\spline_then_smooth\smoothAfterSpline_category32.json"
    output_dir = r"..\fujian\fujian1\pmdarima_single_batch\unit_test_output"
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, "unit_test_output_category32.json")
    
    predict_and_save_inventory(input_file, output_file)
